#!/usr/bin/env python
import rospy
import numpy as np
from geometry_msgs.msg import Twist
from std_msgs.msg import Float64MultiArray
from threading import Lock
from spatialmath.base import eul2r
from roboticstoolbox import DHRobot, RevoluteDH, PrismaticDH
from math import pi
from std_srvs.srv import Empty, EmptyResponse
from robot_solver_pkg.srv import ArmControl, ArmControlResponse


class Project_Arm(DHRobot):
    def __init__(self):
        deg = pi / 180
        mm = 1

        # 定义机械臂的DH参数
        P1 = PrismaticDH(theta=0, a=0., alpha=-pi/2, qlim=[-0.1, 0.1])
        R1 = RevoluteDH(d=0 * mm, a=0., alpha=-pi/2, qlim=[-170 * deg, 170 * deg])
        R2 = RevoluteDH(d=0., a=325.9 * mm, alpha=pi, qlim=[-170 * deg, 170 * deg])
        R3 = RevoluteDH(d=59.6 * mm, a=239.78 * mm, alpha=pi, qlim=[-170 * deg, 170 * deg])
        R4 = RevoluteDH(d=59.53 * mm, a=0., alpha=pi / 2, qlim=[-170 * deg, 170 * deg], offset=pi / 2)
        R5 = RevoluteDH(d=66.14 * mm, a=0., alpha=pi / 2, qlim=[-170 * deg, 170 * deg], offset=-pi / 2)
        R6 = RevoluteDH(d=104.5 * mm, a=0., alpha=0, qlim=[-170 * deg, 170 * deg])

        L = [P1, R1, R2, R3, R4, R5, R6]

        super().__init__(L, name="Project arm", manufacturer="Victor Scheinman", keywords=("dynamics",))

        self.qr = np.zeros(7)
        self.qz = np.zeros(7)

        self.addconfiguration("qr", self.qr)
        self.addconfiguration("qz", self.qz)


class ArmSolverNode:
    def __init__(self):
        # 初始化机械臂模型
        self.arm = Project_Arm()

        # ROS节点初始化
        rospy.init_node('arm_robot_rtbx_solver')

        # 创建服务
        self.service = rospy.Service('arm_control_service', ArmControl, self.handle_request)

        # 状态控制变量
        self.lock = Lock()
        self.vision_active = True
        self.current_joints = self.arm.qz

        rospy.loginfo("逆向运动学解算服务已启动")

    def handle_request(self, req):
        """处理服务请求"""
        if self.lock.acquire(blocking=True):
            try:
                # 提取请求数据
                current_joints = np.array(req.current_joints)
                pose_delta = req.pose_delta

                # 计算目标位姿
                T = self.arm.fkine(current_joints)
                T.t += np.array([pose_delta.linear.x, pose_delta.linear.y, pose_delta.linear.z])
                # Create rotation matrix from euler angles and multiply with current rotation
                delta_R = eul2r([pose_delta.angular.x, pose_delta.angular.y, pose_delta.angular.z])
                # T.R = T.R @ delta_R
                T.R = delta_R

                # 逆向运动学求解
                sol = self.arm.ikine_LM(T, q0=current_joints, joint_limits=True)

                # 返回解算结果
                response = ArmControlResponse()
                response.joint_angles = sol.q.tolist()
                return response
            finally:
                self.lock.release()


if __name__ == '__main__':
    try:
        node = ArmSolverNode()
        rospy.spin()
    except rospy.ROSInterruptException:
        pass